#include <iostream>
#include <opencv2/opencv.hpp>
#include <random>
#include <time.h>

#include "param.h"

const std::string temp_path = "dataset/template/2.jpg";
const std::string output_path = "train/2/";
std::mt19937 rnd(time(0));
constexpr int example_count = 1000; // 数据组数

void srcImgInit(cv::Mat& img) {
  cv::Mat realistic_icon;
  cv::GaussianBlur(img, realistic_icon, cv::Size(5, 5), 0.8); // 轻微模糊边缘

  // 随机调整一些像素值，模拟传感器捕捉到的非纯黑白
  // 增加一些随机的亮度偏移，使图像不是纯粹的 0 或 255
  cv::Mat rand_offset(realistic_icon.size(), realistic_icon.type());
  cv::randn(rand_offset, cv::Scalar(0), cv::Scalar(10)); // 0-10的随机偏移
  cv::add(realistic_icon, rand_offset, realistic_icon);

  // 将图像转为 CV_32F 类型进行后续处理，确保精度
  realistic_icon.convertTo(realistic_icon, CV_32F);
}

// 随机平移
void randomTranslate(cv::Mat &src, cv::Mat &dst) {
  const int MAX_SHIFT_X = src.cols * 0.2;
  const int MAX_SHIFT_Y = src.rows * 0.2;

  int tx = std::uniform_int_distribution<>(-MAX_SHIFT_X, MAX_SHIFT_X)(rnd);
  int ty = std::uniform_int_distribution<>(-MAX_SHIFT_Y, MAX_SHIFT_Y)(rnd);

  cv::Mat M = (cv::Mat_<float>(2, 3) << 1, 0, tx, 0, 1, ty);

  cv::warpAffine(src, dst, M, src.size());
}

// 随机旋转，模拟不同角度
void randomRotate(cv::Mat &src, cv::Mat &dst) {
  cv::Mat temp_src;
  src.convertTo(temp_src, CV_8U); // 临时转换为 CV_8U
  double angle = std::uniform_real_distribution<>(-90, 90)(rnd);
  cv::Point2f center((float)temp_src.cols / 2, (float)temp_src.rows / 2);
  cv::Mat rot_mat = cv::getRotationMatrix2D(center, angle, 1.0);
  cv::Mat temp_dst_8u;
  cv::warpAffine(temp_src, temp_dst_8u, rot_mat, temp_src.size());
  temp_dst_8u.convertTo(dst, CV_32F); // 转回 CV_32F
}

// 随机模糊
void randomBlur(cv::Mat &src, cv::Mat &dst) {
  cv::Mat temp_src;
  src.convertTo(temp_src, CV_8U); // 临时转换为 CV_8U
  int k_size = (std::uniform_int_distribution<>(0, 3)(rnd)) * 2 + 1;
  if (k_size == 1) {
    temp_src.convertTo(dst, CV_32F); // 如果不模糊，也转换为 CV_32F
  } else {
    cv::Mat temp_dst_8u;
    cv::GaussianBlur(temp_src, temp_dst_8u, cv::Size(k_size, k_size), 0);
    temp_dst_8u.convertTo(dst, CV_32F); // 转回 CV_32F
  }
}

// 随机对比度亮度调整
void randomAlphaBeta(cv::Mat &src, cv::Mat &dst) {
  double alpha = std::uniform_real_distribution<>(0.8, 1.2)(rnd);
  int beta = std::uniform_int_distribution<>(-25, 25)(rnd);
  src.convertTo(dst, -1, alpha, beta);
}

// 添加高斯噪声
void addGaussianNoise(cv::Mat& src) {
  // 一组符合 N(mean, std_dev^2) 正态分布的噪点
  double mean = 0.0;
  double std_dev = std::uniform_real_distribution<>(50, 100)(rnd); 
  // 生成噪声矩阵并处理图像
  cv::Mat noise_matrix(src.size(), src.type());
  cv::randn(noise_matrix, cv::Scalar(mean), cv::Scalar(std_dev));
  cv::add(src, noise_matrix, src);
}

int main() {
  cv::Mat img = cv::imread(temp_path);
  if (img.empty()) {
    std::cout << "未找到图片" << std::endl;
    return -1;
  }

  cv::resize(img, img, cv::Size(param::target_height, param::target_width));
  for (int i = 1; i <= example_count; ++i) {
    cv::Mat current_sample;
    randomRotate(img, current_sample);
    randomTranslate(current_sample, current_sample);
    randomAlphaBeta(current_sample, current_sample);
    randomBlur(current_sample, current_sample);
    addGaussianNoise(current_sample);
    
    std::string path = output_path + std::to_string(i) + ".jpg";
    cv::imwrite(path, current_sample);
  }
}